package sjjg.uf;

import java.util.Scanner;

public class UF_Tree_Weight {
    //存储并查集元素及及所在分组(下标表示该元素 而值存储该元素所在分组)
    private int[]eleAndGroup;
    //记录分组数
    private int count;
    //记录树的大小(为了在union时深度降低 提升查找性能)
    private int[]sz;
    //构造函数 初始化并查集
    public UF_Tree_Weight(int N){
        //记录分组数
        count=N;
        //为每个元素设置一个分组
        eleAndGroup=new int[N];
        sz=new int[N];
        for (int i = 0; i < eleAndGroup.length; i++) {
             eleAndGroup[i]=i;
             //初始每棵树即每个组的大小都为1
             sz[i]=1;
        }
    }
    public int count(){
        return count;
    }
    public int find(int ele){
         //找到该组的根节点
        while (true){
            if(eleAndGroup[ele]==ele){
                return ele;
            }
            ele=eleAndGroup[ele];
        }
    }
    public boolean connected(int p,int q){
           return find(p)==find(q);
    }
    public void union(int p,int q){
           if(connected(p,q)){
               return;
           }
        //记录p root
        int pRoot=find(p);
        //记录q root
        int qRoot=find(q);
        //比较两棵树大小(小树合并到大树上提高性能)
        if(sz[pRoot]>sz[qRoot]){
            //将q所在组合并到p组
            eleAndGroup[qRoot]=pRoot;
            //p所在树大小+=q所在树大小
            sz[pRoot]+=sz[qRoot];
        }else {
            //将p所在组合并到q组
            eleAndGroup[pRoot]=qRoot;
            //q所在树大小+=p所在树大小
            sz[qRoot]+=sz[pRoot];
        }
        count--;
    }

    public static void main(String[] args) {
        UF_Tree_Weight uf = new UF_Tree_Weight(7);
        System.out.println("并查集初始化成功，有"+uf.count+"个分组");
        Scanner scanner=new Scanner(System.in);
        while (true){
            System.out.println("请输入第一个元素：");
            int p=scanner.nextInt();
            if(p==212502){
                break;
            }
            System.out.println("请输入第二个元素：");
            int q=scanner.nextInt();
            if(q==212502){
                break;
            }
            if(uf.connected(p,q)){
                System.out.println("已是同一分组...");
                continue;
            }
            System.out.println("正在归并...");
            uf.union(p,q);
            System.out.println("q现所在组为第"+uf.find(q)+"组...");
            System.out.println("现有"+uf.count()+"个分组...");
        }
    }
}